# 设置API Key
import os
os.environ["DASHSCOPE_API_KEY"] = 'sk-4e88cf4db3e14894bafaff606d296610'

# 导入所需的库和模块
from collections import deque
from typing import Dict, List, Optional, Any
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.llms import BaseLLM
from langchain_community.llms import Tongyi  # 使用阿里云的通义千问模型
from langchain.vectorstores.base import VectorStore
from pydantic import BaseModel, Field
from langchain.chains.base import Chain
from langchain_community.vectorstores import FAISS
import faiss
from langchain_community.docstore.in_memory import InMemoryDocstore

print("这是一个一个基于BabyAGI框架的智能任务自动化处理系统，主要用于分解和执行复杂目标（如\"分析北京市今天的气候情况，写出鲜花储存策略\"）。"
      "系统通过不断生成、执行和优化任务来逐步完成目标")
print("""自主任务执行系统：给定一个总体目标（如"分析北京市今天的气候情况，写出鲜花储存策略"），系统能够：
1.自动将大目标分解成小任务
2.按优先级执行任务
3.根据执行结果生成新任务
4.不断优化任务执行顺序
5.存储和利用历史执行结果""")

# 定义嵌入模型
embeddings_model = HuggingFaceEmbeddings(
    model_name="D:/ideaSpace/MyPython/models/bge-small-zh-v1.5",
    model_kwargs={"device": "cpu"},
    encode_kwargs={"normalize_embeddings": True}
)

# 初始化向量存储 - 使用正确的方式创建FAISS索引
def create_vectorstore():
    """创建并返回FAISS向量存储"""
    # 测试嵌入维度
    test_embedding = embeddings_model.embed_query("测试")
    embedding_size = len(test_embedding)
    print(f"嵌入维度: {embedding_size}")

    # 创建FAISS索引
    index = faiss.IndexFlatL2(embedding_size)

    # 创建向量存储
    return FAISS(
        embedding_function=embeddings_model,
        index=index,
        docstore=InMemoryDocstore(),
        index_to_docstore_id={}
    )

# 创建向量存储
vectorstore = create_vectorstore()


# 任务生成链
class TaskCreationChain(LLMChain):
    """负责生成任务的链"""
    @classmethod
    def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain:
        """从LLM获取响应解析器"""
        task_creation_template = (
            "你是一个任务创建AI，使用执行代理的结果来创建新任务。"
            "目标：{objective}，"
            "上一个完成的任务结果：{result}。"
            "这个结果基于任务描述：{task_description}。"
            "这些是未完成的任务：{incomplete_tasks}。"
            "基于结果，创建新的任务，不要与未完成的任务重叠。"
            "以数组格式返回任务，每个任务一行。"
        )
        prompt = PromptTemplate(
            template=task_creation_template,
            input_variables=[
                "result",
                "task_description",
                "incomplete_tasks",
                "objective",
            ],
        )
        return cls(prompt=prompt, llm=llm, verbose=verbose)


# 任务优先级链
class TaskPrioritizationChain(LLMChain):
    """负责任务优先级排序的链"""
    @classmethod
    def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain:
        """从大模型获取响应解析器"""
        task_prioritization_template = (
            "你是一个任务优先级排序AI，负责清理格式和重新排序以下任务：{task_names}。"
            "考虑你的团队的最终目标：{objective}。"
            "不要删除任何任务。以编号列表格式返回结果，例如："
            "1. 第一个任务"
            "2. 第二个任务"
            "任务列表从数字 {next_task_id} 开始。"
        )
        prompt = PromptTemplate(
            template=task_prioritization_template,
            input_variables=["task_names", "next_task_id", "objective"],
        )
        return cls(prompt=prompt, llm=llm, verbose=verbose)


# 任务执行链
class ExecutionChain(LLMChain):
    """负责执行任务的链"""

    @classmethod
    def from_llm(cls, llm: BaseLLM, verbose: bool = True) -> LLMChain:
        """从大模型获取响应解析器"""
        execution_template = (
            "你是一个AI，基于以下目标执行一个任务：{objective}。"
            "考虑这些先前完成的任务：{context}。"
            "你的任务：{task}。"
            "请用中文详细回答："
        )
        prompt = PromptTemplate(
            template=execution_template,
            input_variables=["objective", "context", "task"],
        )
        return cls(prompt=prompt, llm=llm, verbose=verbose)


# 获取下一个任务
def get_next_task(
        task_creation_chain: LLMChain,
        result: str,
        task_description: str,
        task_list: List[str],
        objective: str,
) -> List[Dict]:
    """Get the next task."""
    incomplete_tasks = ", ".join(task_list)
    response = task_creation_chain.run(
        result=result,
        task_description=task_description,
        incomplete_tasks=incomplete_tasks,
        objective=objective,
    )

    # 解析响应，提取任务
    new_tasks = []
    for line in response.split('\n'):
        line = line.strip()
        if line and not line.startswith(('#', '-', '1.', '2.', '3.', '4.', '5.')):
            # 移除可能的前缀
            if '. ' in line:
                line = line.split('. ', 1)[1]
            new_tasks.append({"task_name": line})

    return new_tasks

# 设置任务优先级
def prioritize_tasks(
        task_prioritization_chain: LLMChain,
        this_task_id: int,
        task_list: List[Dict],
        objective: str,
) -> List[Dict]:
    """Prioritize tasks."""
    task_names = [t["task_name"] for t in task_list]
    next_task_id = int(this_task_id) + 1

    try:
        response = task_prioritization_chain.run(
            task_names=task_names, next_task_id=next_task_id, objective=objective
        )
    except Exception as e:
        print(f"优先级排序失败，使用原始顺序: {e}")
        return task_list

    prioritized_task_list = []
    for line in response.split('\n'):
        line = line.strip()
        if line and any(line.startswith(f"{i}.") for i in range(1, 10)):
            parts = line.split('. ', 1)
            if len(parts) == 2:
                task_id = parts[0].strip()
                task_name = parts[1].strip()
                prioritized_task_list.append({"task_id": task_id, "task_name": task_name})

    # 如果解析失败，返回原始任务列表
    if not prioritized_task_list:
        print("无法解析优先级响应，使用原始顺序")
        for i, task in enumerate(task_list, 1):
            prioritized_task_list.append({"task_id": str(i), "task_name": task["task_name"]})

    return prioritized_task_list

# 获取头部任务
def _get_top_tasks(vectorstore, query: str, k: int) -> List[str]:
    """Get the top k tasks based on the query."""
    try:
        results = vectorstore.similarity_search(query, k=k)
        if not results:
            return []
        return [str(item.metadata["task"]) for item in results]
    except Exception as e:
        print(f"相似性搜索错误: {e}")
        return []

# 执行任务
def execute_task(
        vectorstore, execution_chain: LLMChain, objective: str, task: str, k: int = 3
) -> str:
    """Execute a task."""
    context = _get_top_tasks(vectorstore, query=objective, k=k)
    print(f"检索到的上下文: {context}")

    # 使用正确的invoke方法调用方式
    result = execution_chain.invoke({
        "objective": objective,
        "context": context,
        "task": task
    })
    return result["text"]  # 从返回结果中提取文本


# BabyAGI 主类
class BabyAGI(Chain, BaseModel):
    """BabyAGI代理的控制器模型"""

    task_list: deque = Field(default_factory=deque)
    task_creation_chain: TaskCreationChain = Field(...)
    task_prioritization_chain: TaskPrioritizationChain = Field(...)
    execution_chain: ExecutionChain = Field(...)
    task_id_counter: int = Field(1)
    vectorstore: VectorStore = Field(...)
    max_iterations: Optional[int] = None

    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True

    def add_task(self, task: Dict):
        self.task_list.append(task)

    def print_task_list(self):
        print("\033[95m\033[1m" + "\n*****任务列表*****\n" + "\033[0m\033[0m")
        for t in self.task_list:
            print(f"{t['task_id']}: {t['task_name']}")

    def print_next_task(self, task: Dict):
        print("\033[92m\033[1m" + "\n*****下一个任务*****\n" + "\033[0m\033[0m")
        print(f"{task['task_id']}: {task['task_name']}")

    def print_task_result(self, result: str):
        print("\033[93m\033[1m" + "\n*****任务结果*****\n" + "\033[0m\033[0m")
        print(result)

    @property
    def input_keys(self) -> List[str]:
        return ["objective"]

    @property
    def output_keys(self) -> List[str]:
        return []

    def _call(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Run the agent."""
        objective = inputs["objective"]
        first_task = inputs.get("first_task", "制定任务清单")
        self.add_task({"task_id": "1", "task_name": first_task})
        num_iters = 0

        # 添加初始任务到向量存储
        try:
            self.vectorstore.add_texts(
                texts=[first_task],
                metadatas=[{"task": first_task}],
                ids=["initial_task"]
            )
        except Exception as e:
            print(f"向量存储初始化错误: {e}")

        while self.task_list and num_iters < (self.max_iterations or 5):
            self.print_task_list()

            # Step 1: Pull the first task
            task = self.task_list.popleft()
            self.print_next_task(task)

            # Step 2: Execute the task
            try:
                result = execute_task(
                    self.vectorstore, self.execution_chain, objective, task["task_name"]
                )
                this_task_id = int(task["task_id"])
                self.print_task_result(result)

                # Step 3: Store the result in vectorstore
                try:
                    result_id = f"result_{task['task_id']}_{num_iters}"
                    self.vectorstore.add_texts(
                        texts=[result],
                        metadatas=[{"task": task["task_name"]}],
                        ids=[result_id],
                    )
                except Exception as e:
                    print(f"存储结果到向量数据库错误: {e}")

                # Step 4: Create new tasks
                try:
                    new_tasks = get_next_task(
                        self.task_creation_chain,
                        result,
                        task["task_name"],
                        [t["task_name"] for t in self.task_list],
                        objective,
                    )
                    for new_task in new_tasks:
                        self.task_id_counter += 1
                        new_task.update({"task_id": str(self.task_id_counter)})
                        self.add_task(new_task)
                except Exception as e:
                    print(f"创建新任务错误: {e}")

                # Step 5: Reprioritize task list
                if self.task_list:
                    try:
                        prioritized_tasks = prioritize_tasks(
                            self.task_prioritization_chain,
                            this_task_id,
                            list(self.task_list),
                            objective,
                        )
                        self.task_list = deque(prioritized_tasks)
                    except Exception as e:
                        print(f"任务优先级排序错误: {e}")

            except Exception as e:
                print(f"执行任务错误: {e}")
                # 如果任务执行失败，重新添加任务
                self.add_task(task)

            num_iters += 1
            print(f"\033[94m\033[1m迭代次数: {num_iters}\033[0m\033[0m")

        print("\033[91m\033[1m" + "\n*****任务结束*****\n" + "\033[0m\033[0m")
        return {}

    @classmethod
    def from_llm(
            cls, llm: BaseLLM, vectorstore: VectorStore, verbose: bool = False, **kwargs
    ) -> "BabyAGI":
        """Initialize the BabyAGI Controller."""
        task_creation_chain = TaskCreationChain.from_llm(llm, verbose=verbose)
        task_prioritization_chain = TaskPrioritizationChain.from_llm(
            llm, verbose=verbose
        )
        execution_chain = ExecutionChain.from_llm(llm, verbose=verbose)
        return cls(
            task_creation_chain=task_creation_chain,
            task_prioritization_chain=task_prioritization_chain,
            execution_chain=execution_chain,
            vectorstore=vectorstore,
            **kwargs,
        )


# 主函数执行部分
if __name__ == "__main__":
    OBJECTIVE = "分析一下北京市今天的气候情况，写出鲜花储存策略。"

    # 使用阿里云通义千问模型
    llm = Tongyi(
        dashscope_api_key="sk-4e88cf4db3e14894bafaff606d296610",
        model_name="qwen-turbo",  # 可以使用 qwen-turbo, qwen-plus, qwen-max 等
        temperature=0.1,
        top_p=0.8
    )

    verbose = True
    max_iterations = 3

    baby_agi = BabyAGI.from_llm(
        llm=llm,
        vectorstore=vectorstore,
        verbose=verbose,
        max_iterations=max_iterations
    )

    try:
        result = baby_agi.invoke({"objective": OBJECTIVE})
        print("\033[92m\033[1m" + "\n*****最终完成*****\n" + "\033[0m\033[0m")
    except Exception as e:
        print(f"运行过程中出现错误: {e}")
        print(f"错误类型: {type(e).__name__}")
        import traceback
        traceback.print_exc()

"""工作流程
1.初始化：创建初始任务（"制定任务清单"）
2.执行循环：
  取出最高优先级任务
  执行任务并获取结果
  存储结果到向量数据库
  根据结果生成新任务
  重新排序所有任务
3.终止条件：达到最大迭代次数或任务完成

技术特点
    自主性：无需人工干预，自动分解和执行任务
    适应性：根据执行结果动态调整任务计划
    记忆性：利用历史执行结果优化后续任务
    可扩展性：可以处理各种类型的复杂任务
这种系统非常适合需要多步骤推理、知识积累和自主决策的复杂任务处理场景！
"""

"""核心组件作用
1. BabyAGI类 - 大脑中枢
任务管理：维护任务队列，管理任务生命周期

协调控制：协调任务创建、执行、优先级排序的整个流程

迭代循环：控制任务的执行迭代过程

状态维护：跟踪任务ID计数器和执行状态

2. 三个核心链（Chains）
🎨 TaskCreationChain（任务生成链）
功能：根据当前任务执行结果，智能生成新的子任务

输入：目标、上一个任务结果、未完成任务列表

输出：新的任务列表

示例：执行"获取天气数据"后，可能生成"分析温度对鲜花的影响"等新任务

📊 TaskPrioritizationChain（任务优先级链）
功能：对任务列表进行智能排序和优先级调整

输入：所有待处理任务、总体目标

输出：按优先级排序的任务列表

示例：将"查询实时天气"排在"分析历史数据"前面

⚡ ExecutionChain（任务执行链）
功能：实际执行具体任务

输入：目标任务、相关上下文、具体任务描述

输出：任务执行结果

示例：执行"查询北京今日温度"返回"今日北京气温25°C"

3. 向量数据库（FAISS）
记忆功能：存储所有已完成任务的结果

相关性检索：为新任务提供相关的历史上下文

知识积累：系统通过不断执行任务积累知识
"""